from fastapi import FastAPI, WebSocket
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import uvicorn
import asyncio
import random
import time
import json
import numpy as np

app = FastAPI()

# 存储活跃连接
active_connections = set()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    active_connections.add(websocket)
    print("🎯 客户端连接成功")
    
    try:
        # 启动数据生成任务
        data_task = asyncio.create_task(send_waveform_data(websocket))
        
        while True:
            # 接收消息
            data = await websocket.receive_text()
            print(f"📨 收到消息: {data}")
            
            # 处理前端发送的pingclireq
            if data == "pingclireq":
                print("✅ 收到pingclireq，发送pongsvrresp")
                await websocket.send_text("pongsvrresp")
            
            # 处理前端响应的pongcliresp
            elif data == "pongcliresp":
                print("✅ 收到pongcliresp，连接正常")
                await websocket.send_text("ack")
            
            # 处理其他消息
            else:
                print(f"📝 收到其他消息: {data}")
                await websocket.send_text(f"收到: {data}")
                
    except Exception as e:
        print(f"❌ 连接错误: {e}")
    finally:
        data_task.cancel()
        active_connections.remove(websocket)
        print("🔴 客户端断开连接")

# 生成更真实的波形数据
def generate_waveform_data(num_points):
    max_value = 2**16 - 1  # uint16_t最大值 (0-65535)
    
    # 生成基波 (1Hz)
    t = np.linspace(0, 20 * 60, num_points)  # 20分钟时间轴
    base_wave = np.sin(2 * np.pi * 1 * t) * 0.4 * max_value
    
    # 添加谐波 (3Hz, 5Hz)
    harmonic1 = np.sin(2 * np.pi * 3 * t) * 0.2 * max_value
    harmonic2 = np.sin(2 * np.pi * 5 * t) * 0.1 * max_value
    
    # 添加随机噪声
    noise = np.random.normal(0, 0.05 * max_value, num_points)
    
    # 合成波形
    waveform = base_wave + harmonic1 + harmonic2 + noise
    
    # 转换为uint16范围
    waveform = np.clip(waveform, 0, max_value)
    return waveform.astype(int).tolist()

# 生成并发送波形数据
async def send_waveform_data(websocket):
    sample_rate = 1000  # 1000Hz采样率
    interval = 0.02     # 20ms间隔
    bits = 16           # 16位精度
    points_per_chunk = 20  # 每次发送20个点
    
    # 预生成20分钟数据 (1000Hz * 60 * 20 = 1,200,000点)
    total_points = sample_rate * 60 * 20
    waveform_data = generate_waveform_data(total_points)
    data_index = 0
    
    try:
        while True:
            # 获取当前数据块
            chunk = waveform_data[data_index:data_index + points_per_chunk]
            data_index = (data_index + points_per_chunk) % total_points
            
            # 构建数据包
            data = {
                "type": "waveform",
                "data": chunk,
                "sample_rate": sample_rate,
                "bits": bits,
                "timestamp": time.time(),
                "total_points": total_points,
                "current_index": data_index
            }
            
            # 发送JSON数据
            await websocket.send_json(data)
            await asyncio.sleep(interval)
            
    except asyncio.CancelledError:
        print("波形数据发送任务已取消")
    except Exception as e:
        print(f"❌ 数据发送错误: {e}")

# 定期发送pingsvrreq的任务
async def send_pingsvrreq():
    while True:
        await asyncio.sleep(30)  # 每30秒发送一次pingsvrreq
        if active_connections:
            print(f"🔄 发送pingsvrreq到 {len(active_connections)} 个客户端")
            for connection in list(active_connections):
                try:
                    await connection.send_text("pingsvrreq")
                    print("✅ pingsvrreq发送成功")
                except Exception as e:
                    print(f"❌ pingsvrreq发送失败: {e}")
                    active_connections.discard(connection)

@app.get("/")
async def get():
    with open("index.html") as f:
        return HTMLResponse(f.read())

async def main():
    # 启动pingsvrreq发送任务
    ping_task = asyncio.create_task(send_pingsvrreq())
    
    # 启动Web服务器
    config = uvicorn.Config(app, host="0.0.0.0", port=8000)
    server = uvicorn.Server(config)
    
    try:
        await server.serve()
    finally:
        ping_task.cancel()
        try:
            await ping_task
        except asyncio.CancelledError:
            print("pingsvrreq任务已取消")

if __name__ == "__main__":
    asyncio.run(main())